#!/usr/bin/env python2
# -*- coding: utf-8 -*-

from __future__ import print_function
from triton import TritonContext, ARCH, MODE, AST_REPRESENTATION, Instruction, OPERAND
from triton import *
import sys
import os
import lief

# 加载 elf 文件到内存


INPUT_ADDR = 0x100000

RBP_ADDR = 0x600000
RSP_ADDR = RBP_ADDR - 0x200000


def loadBinary(ctx, path):
    binary = lief.parse(path)
    phdrs = binary.segments
    for phdr in phdrs:
        size = phdr.physical_size
        vaddr = phdr.virtual_address
        print('[+] Loading 0x%06x - 0x%06x' % (vaddr, vaddr+size))
        ctx.setConcreteMemoryAreaValue(vaddr, phdr.content)
    return

if __name__ == '__main__':
    ctx = TritonContext()
    ctx.setArchitecture(ARCH.X86_64)
    ctx.enableMode(MODE.ALIGNED_MEMORY, True)

    loadBinary(ctx, os.path.join(os.path.dirname(__file__), 'crackme_xor'))
    ctx.setAstRepresentationMode(AST_REPRESENTATION.PYTHON)

    pc = 0x0400556

    # 参数是输入字符串的指针
    ctx.setConcreteRegisterValue(ctx.registers.rdi, INPUT_ADDR)

    # 设置栈的值
    ctx.setConcreteRegisterValue(ctx.registers.rsp, RSP_ADDR)
    ctx.setConcreteRegisterValue(ctx.registers.rbp, RBP_ADDR)

    # ctx.taintRegister(ctx.registers.rdi)

    input = "elite\x00"
    ctx.setConcreteMemoryAreaValue(INPUT_ADDR, input)
    ctx.taintMemory(MemoryAccess(INPUT_ADDR, 8))

    while pc != 0x4005B1:
        # Build an instruction
        inst = Instruction()
        opcode = ctx.getConcreteMemoryAreaValue(pc, 16)
        inst.setOpcode(opcode)
        inst.setAddress(pc)

        # 执行指令
        ctx.processing(inst)

        if inst.isTainted():
            # print('[tainted] %s' % (str(inst)))

            if inst.isMemoryRead():
                for op in inst.getOperands():
                    if op.getType() == OPERAND.MEM:
                        print("read:0x{:08x}, size:{}".format(
                            op.getAddress(), op.getSize()))

            if inst.isMemoryWrite():
                for op in inst.getOperands():
                    if op.getType() == OPERAND.MEM:
                        print("write:0x{:08x}, size:{}".format(
                            op.getAddress(), op.getSize()))

        # 取出下一条指令的地址
        pc = ctx.getConcreteRegisterValue(ctx.registers.rip)

    sys.exit(0)
